﻿using MS.Win32;
using System;
using System.Runtime.InteropServices;
using System.ServiceProcess;

namespace ThinkPadDockHelper
{
    public partial class TPDH : ServiceBase
    {
        private NativeMethods.SERVICE_STATUS status = new NativeMethods.SERVICE_STATUS();
        private IntPtr statusHandle;
        private NativeMethods.ServiceControlCallbackEx commandCallbackEx;
        private IntPtr hDeviceNotification;
        private IntPtr hPowerSettingNotification;
        private bool lidOpened;

        private static readonly string TPDockDeviceId = "VID_17E9&PID_4302";

        public TPDH()
        {
            InitializeComponent();

            status.serviceType = NativeMethods.SERVICE_WIN32;
            status.currentState = NativeMethods.SERVICE_START_PENDING;
            status.controlsAccepted =
                NativeMethods.SERVICE_ACCEPT_STOP |
                NativeMethods.SERVICE_ACCEPT_SHUTDOWN |
                NativeMethods.SERVICE_ACCEPT_HARDWAREPROFILECHANGE |
                NativeMethods.SERVICE_ACCEPT_POWEREVENT;
            status.win32ExitCode = 0;
            status.serviceSpecificExitCode = 0;
            status.checkPoint = 0;
            status.waitHint = 0;

            this.commandCallbackEx = new NativeMethods.ServiceControlCallbackEx(this.ServiceCommandCallbackEx);
        }

        protected override void OnStart(string[] args)
        {
            status.currentState = NativeMethods.SERVICE_RUNNING;

            statusHandle = NativeMethods.RegisterServiceCtrlHandlerEx(ServiceName, this.commandCallbackEx, IntPtr.Zero);

            NativeMethods.DEV_BROADCAST_DEVICEINTERFACE notificationFilter = new NativeMethods.DEV_BROADCAST_DEVICEINTERFACE();
            var size = Marshal.SizeOf(notificationFilter);
            notificationFilter.dbcc_size = (uint)size;
            notificationFilter.dbcc_devicetype = NativeMethods.DBT_DEVTYP_DEVICEINTERFACE;
            notificationFilter.dbcc_classguid = NativeMethods.GUID_DEVINTERFACE_USB_DEVICE;
            var notificationFilterData = Marshal.AllocHGlobal(size);
            Marshal.StructureToPtr(notificationFilter, notificationFilterData, true);
            hDeviceNotification = UnsafeNativeMethods.RegisterDeviceNotification(this.ServiceHandle, notificationFilterData, NativeMethods.DEVICE_NOTIFY_SERVICE_HANDLE);
            Marshal.FreeHGlobal(notificationFilterData);

            var powerSettingGuid = NativeMethods.GUID_LIDSWITCH_STATE_CHANGE;
            hPowerSettingNotification = UnsafeNativeMethods.RegisterPowerSettingNotification(this.ServiceHandle, ref powerSettingGuid, NativeMethods.DEVICE_NOTIFY_SERVICE_HANDLE);
        }

        protected override void OnStop()
        {
            if (hDeviceNotification != IntPtr.Zero)
            {
                UnsafeNativeMethods.UnregisterDeviceNotification(hDeviceNotification);
                hDeviceNotification = IntPtr.Zero;
            }

            if (hPowerSettingNotification != IntPtr.Zero)
            {
                UnsafeNativeMethods.UnregisterPowerSettingNotification(hPowerSettingNotification);
                hPowerSettingNotification = IntPtr.Zero;
            }
        }

        private delegate void DeferredHandlerDelegate();
        private delegate void DeferredHandlerDelegateAdvanced(int eventType, IntPtr eventData);

        private void DeferredPowerEvent(int eventType, IntPtr eventData)
        {
            switch (eventType)
            {
                case NativeMethods.PBT_POWERSETTINGCHANGE:
                    var ps = (NativeMethods.POWERBROADCAST_SETTING)Marshal.PtrToStructure(eventData, typeof(NativeMethods.POWERBROADCAST_SETTING));
                    if (ps.PowerSetting == NativeMethods.GUID_LIDSWITCH_STATE_CHANGE)
                        lidOpened = Convert.ToBoolean(ps.Data);

                    break;

                default:
                    break;
            }
        }

        private void DeferredDeviceEvent(int eventType, IntPtr eventData)
        {
            var dbhdr = (NativeMethods.DEV_BROADCAST_HDR)Marshal.PtrToStructure(eventData, typeof(NativeMethods.DEV_BROADCAST_HDR));
            if (dbhdr.dbch_devicetype == NativeMethods.DBT_DEVTYP_DEVICEINTERFACE)
            {
                var dbdi = (NativeMethods.DEV_BROADCAST_DEVICEINTERFACE)Marshal.PtrToStructure(eventData, typeof(NativeMethods.DEV_BROADCAST_DEVICEINTERFACE));
                if (dbdi.dbcc_name.Contains(TPDockDeviceId))
                {
                    uint rc = 0;

                    // Get the active power scheme
                    Guid activePolicyGuid = Guid.Empty;
                    rc = PowerUtil.PowerGetActiveScheme(ref activePolicyGuid);

                    switch (eventType)
                    {
                        case NativeMethods.DBT_DEVICEARRIVAL:
                            // Read the lid close action value index
                            var currentACValueIndex = default(uint);
                            rc = PowerUtil.PowerReadACValueIndex(ref activePolicyGuid, ref currentACValueIndex);

                            var currentDCValueIndex = default(uint);
                            rc = PowerUtil.PowerReadDCValueIndex(ref activePolicyGuid, ref currentDCValueIndex);

                            // Update the stored value index so we can revert to it when undocking
                            RegistryUtil.SetUndockedLidCloseActionACValueIndex(currentACValueIndex);
                            RegistryUtil.SetUndockedLidCloseActionDCValueIndex(currentDCValueIndex);

                            // Update the lid close action value index to do nothing
                            rc = PowerUtil.PowerWriteACValueIndex(ref activePolicyGuid, NativeMethods.POWERBUTTON_ACTION_INDEX_NOTHING);
                            rc = PowerUtil.PowerWriteDCValueIndex(ref activePolicyGuid, NativeMethods.POWERBUTTON_ACTION_INDEX_NOTHING);

                            // Set the active power scheme to apply the changes
                            rc = PowerUtil.PowerSetActiveScheme(ref activePolicyGuid);

                            break;

                        case NativeMethods.DBT_DEVICEREMOVECOMPLETE:
                            // Retrieve the stored value index so we can revert to it
                            var previousACValueIndex = RegistryUtil.GetUndockedLidCloseActionACValueIndex();
                            var previousDCValueIndex = RegistryUtil.GetUndockedLidCloseActionDCValueIndex();

                            // Update the lid close action value index to the stored value index
                            rc = PowerUtil.PowerWriteACValueIndex(ref activePolicyGuid, previousACValueIndex);
                            rc = PowerUtil.PowerWriteDCValueIndex(ref activePolicyGuid, previousDCValueIndex);

                            // Set the active power scheme to apply the changes
                            rc = PowerUtil.PowerSetActiveScheme(ref activePolicyGuid);

                            // If the dock was disconnected while the lid was closed, then we suspend the system
                            if (!lidOpened)
                                UnsafeNativeMethods.SetSuspendState(false, false, true);

                            break;
                    }
                }
            }
        }

        private unsafe void DeferredStop()
        {
            fixed (NativeMethods.SERVICE_STATUS* pStatus = &status)
            {
                var previousState = status.currentState;

                status.currentState = NativeMethods.SERVICE_STOP_PENDING;
                NativeMethods.SetServiceStatus(statusHandle, pStatus);

                try
                {
                    OnStop();
                    //TODO:logging
                    status.currentState = NativeMethods.SERVICE_STOPPED;
                    NativeMethods.SetServiceStatus(statusHandle, pStatus);

                    try
                    {
                        AppDomain.Unload(AppDomain.CurrentDomain);
                    }
                    catch (CannotUnloadAppDomainException e)
                    {
                        //TODO:logging
                    }
                }
                catch (Exception e)
                {
                    status.currentState = previousState;
                    NativeMethods.SetServiceStatus(statusHandle, pStatus);
                    //TODO:logging
                    throw;
                }
            }
        }

        private unsafe void DeferredShutdown()
        {
            try
            {
                OnShutdown();
                //TODO:logging

                if (status.currentState == NativeMethods.SERVICE_RUNNING)
                {
                    fixed (NativeMethods.SERVICE_STATUS* pStatus = &status)
                    {
                        status.currentState = NativeMethods.SERVICE_STOPPED;
                        NativeMethods.SetServiceStatus(statusHandle, pStatus);

                        try
                        {
                            AppDomain.Unload(AppDomain.CurrentDomain);
                        }
                        catch (CannotUnloadAppDomainException e)
                        {
                            //TODO:logging
                        }
                    }
                }
            }
            catch (Exception e)
            {
                //TODO:logging
                throw;
            }
        }

        private int ServiceCommandCallbackEx(int command, int eventType, IntPtr eventData, IntPtr eventContext)
        {
            var result = NativeMethods.NO_ERROR;

            switch (command)
            {
                case NativeMethods.SERVICE_CONTROL_POWEREVENT:
                    DeferredHandlerDelegateAdvanced powerDelegate = new DeferredHandlerDelegateAdvanced(DeferredPowerEvent);
                    powerDelegate.BeginInvoke(eventType, eventData, null, null);

                    break;

                case NativeMethods.SERVICE_CONTROL_DEVICEEVENT:
                    DeferredHandlerDelegateAdvanced deviceDelegate = new DeferredHandlerDelegateAdvanced(DeferredDeviceEvent);
                    deviceDelegate.BeginInvoke(eventType, eventData, null, null);

                    break;

                default:
                    ServiceCommandCallback(command);
                    break;
            }

            return result;
        }

        private unsafe void ServiceCommandCallback(int command)
        {
            fixed (NativeMethods.SERVICE_STATUS* pStatus = &status)
            {
                if (command == NativeMethods.SERVICE_CONTROL_INTERROGATE)
                    NativeMethods.SetServiceStatus(statusHandle, pStatus);
                else if (status.currentState != NativeMethods.SERVICE_STOP_PENDING)
                {
                    switch (command)
                    {
                        case NativeMethods.SERVICE_CONTROL_STOP:
                            var previousState = status.currentState;
                            if (status.currentState == NativeMethods.SERVICE_RUNNING)
                            {
                                status.currentState = NativeMethods.SERVICE_STOP_PENDING;
                                NativeMethods.SetServiceStatus(statusHandle, pStatus);
                                // Set our copy of the state back to the previous so that the deferred stop routine
                                // can also save the previous state.
                                status.currentState = previousState;

                                DeferredHandlerDelegate stopDelegate = new DeferredHandlerDelegate(DeferredStop);
                                stopDelegate.BeginInvoke(null, null);
                            }

                            break;

                        case NativeMethods.SERVICE_CONTROL_SHUTDOWN:
                            DeferredHandlerDelegate shutdownDelegate = new DeferredHandlerDelegate(DeferredShutdown);

                            shutdownDelegate.BeginInvoke(null, null);
                            break;
                    }
                }
            }
        }
    }
}
